import tkinter as tk
from tkinter import messagebox, filedialog
import os
import socket
import threading
from http.server import HTTPServer, BaseHTTPRequestHandler, SimpleHTTPRequestHandler

class CustomHandler(SimpleHTTPRequestHandler):
    def __init__(self, *args, directory="received_files", **kwargs):
        super().__init__(*args, directory=directory, **kwargs)

    def do_GET(self):
        """Handle GET requests"""
        # 如果请求的不是目录列表
        if not self.path == '/':
            try:
                # 解码 URL 路径
                file_path = os.path.join(self.directory, os.path.basename(self.path))
                if os.path.exists(file_path) and os.path.isfile(file_path):
                    # 设置响应头，强制下载
                    self.send_response(200)
                    self.send_header('Content-Type', 'application/octet-stream')
                    self.send_header('Content-Disposition', f'attachment; filename="{os.path.basename(file_path)}"')
                    self.send_header('Content-Length', str(os.path.getsize(file_path)))
                    self.end_headers()
                    
                    # 发送文件内容
                    with open(file_path, 'rb') as f:
                        self.wfile.write(f.read())
                    return
            except Exception as e:
                self.send_error(404, str(e))
                return
        
        # 如果是目录列表请求，使用原来的处理方法
        return super().do_GET()

    def list_directory(self, path):
        """Override to create custom directory listing page with download buttons"""
        try:
            list = os.listdir(path)
        except OSError:
            self.send_error(404, "No permission to list directory")
            return None
        
        # Create HTML page with bootstrap styling
        r = []
        r.append('<!DOCTYPE HTML>')
        r.append('<html>\n<head>')
        r.append('<meta charset="utf-8">')
        r.append('<title>文件下载</title>')
        r.append('<link href="https://cdn.jsdelivr.net/npm/bootstrap@5.1.3/dist/css/bootstrap.min.css" rel="stylesheet">')
        r.append('</head>')
        r.append('<body class="container mt-5">')
        r.append('<h2>可下载文件列表</h2>')
        r.append('<div class="list-group">')
        
        for name in list:
            fullname = os.path.join(path, name)
            displayname = linkname = name
            if os.path.isfile(fullname):
                r.append(f'<div class="list-group-item d-flex justify-content-between align-items-center">')
                r.append(f'<span>{displayname}</span>')
                r.append(f'<a href="{linkname}" class="btn btn-primary btn-sm">下载</a>')
                r.append('</div>')
        
        r.append('</div>')
        r.append('</body>\n</html>\n')
        
        encoded = '\n'.join(r).encode('utf-8')
        self.send_response(200)
        self.send_header("Content-type", "text/html; charset=utf-8")
        self.send_header("Content-Length", str(len(encoded)))
        self.end_headers()
        self.wfile.write(encoded)
        return None

class transfer:
    def __init__(self, manager):
        self.manager = manager
        self.root = manager.root
        self.server = None
        self.http_server = None
        self.is_server_running = False
        self.server_port = 5000
        self.http_port = 8000
        
        # Create received_files directory if it doesn't exist
        os.makedirs("received_files", exist_ok=True)

    def show(self):
        tk.Label(self.root, text="局域网文件传输工具", font=("Arial", 16)).pack(pady=20)
        
        # 添加服务器控制按钮
        self.server_btn = tk.Button(self.root, text="启动服务器", command=self.toggle_server)
        self.server_btn.pack(pady=10)
        
        # 添加文件选择按钮
        tk.Button(self.root, text="选择文件", command=self.select_file).pack(pady=10)
        
        # 状态显示标签
        self.status_label = tk.Label(self.root, text="服务器状态: 未启动")
        self.status_label.pack(pady=10)
        
        # IP地址显示
        self.ip_label = tk.Label(self.root, text=f"本机IP: {self.get_local_ip() + ':' + str(self.http_port)}")
        self.ip_label.pack(pady=10)
        
        tk.Button(self.root, text="Back to Main", command=self.manager.show_main_page).pack(pady=10)

    def toggle_server(self):
        if not self.is_server_running:
            self.start_server()
        else:
            self.stop_server()

    def start_server(self):
        import socket
        import threading
        
        def server_thread():
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.bind(('0.0.0.0', self.server_port))
                s.listen()
                self.server = s
                self.status_label.config(text="服务器状态: 运行中")
                while self.is_server_running:
                    try:
                        conn, addr = s.accept()
                        threading.Thread(target=self.handle_client, args=(conn, addr)).start()
                    except:
                        break
        
        # Add HTTP server
        def http_server_thread():
            handler = CustomHandler
            self.http_server = HTTPServer(('0.0.0.0', self.http_port), handler)
            while self.is_server_running:
                self.http_server.handle_request()

        self.is_server_running = True
        self.server_btn.config(text="停止服务器")
        threading.Thread(target=server_thread, daemon=True).start()
        threading.Thread(target=http_server_thread, daemon=True).start()
        self.status_label.config(text=f"服务器状态: 运行中\n文件下载地址: http://{self.get_local_ip()}:{self.http_port}")

    def stop_server(self):
        self.is_server_running = False
        if self.server:
            self.server.close()
        if self.http_server:
            self.http_server.server_close()
        self.server_btn.config(text="启动服务器")
        self.status_label.config(text="服务器状态: 已停止")

    def handle_client(self, conn, addr):
        import os
        try:
            filename = conn.recv(1024).decode()
            filesize = int(conn.recv(1024).decode())
            
            with open(os.path.join("received_files", filename), 'wb') as f:
                remaining = filesize
                while remaining:
                    data = conn.recv(min(1024, remaining))
                    if not data:
                        break
                    f.write(data)
                    remaining -= len(data)
        finally:
            conn.close()

    def select_file(self):
        from tkinter import filedialog
        filename = filedialog.askopenfilename()
        if filename:
            try:
                import shutil
                # 直接将选择的文件复制到received_files目录
                dest_path = os.path.join("received_files", os.path.basename(filename))
                shutil.copy2(filename, dest_path)
                tk.messagebox.showinfo("成功", "文件已添加到共享目录！")
            except Exception as e:
                tk.messagebox.showerror("错误", f"文件复制失败: {str(e)}")

    def get_local_ip(self):
        import socket
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect(("8.8.8.8", 80))
            ip = s.getsockname()[0]
            s.close()
            return ip
        except:
            return "无法获取IP"
